# QUIZ
# 
# In the seir_model function below, 
# solve the SEIR model using the 
# Forward Euler Method.

import math
import numpy
import matplotlib.pyplot

h = 0.5 # days
transmission_coeff = 5e-9 # 1 / day person
latency_time = 1. # days
infectious_time = 5. # days

end_time = 60.0 # days
num_steps = int(end_time / h)
times = h * numpy.array(range(num_steps + 1))

def susceptible_derivative(t, i, s):
  return -1 * transmission_coeff * i[t] * s[t] 

def exposed_derivative(t, i, s, e):
  return transmission_coeff * i[t] * s[t] - (1/latency_time) * e[t]

def infectious_derivative(t, i, e):
  return (1/latency_time) * e[t] - (1/infectious_time)*i[t]

def recovered_derivative(t, i):
  return (1/infectious_time) * i[t]

def seir_model():

    s = numpy.zeros(num_steps + 1)
    e = numpy.zeros(num_steps + 1)
    i = numpy.zeros(num_steps + 1)
    r = numpy.zeros(num_steps + 1)

    s[0] = 1e8 - 1e6 - 1e5
    e[0] = 0.
    i[0] = 1e5
    r[0] = 1e6

    for step in range(num_steps):
        s[step+1] = s[step] + h * susceptible_derivative(step, i, s)
        e[step+1] = e[step] + h * exposed_derivative(step, i, s, e)
        i[step+1] = i[step] + h * infectious_derivative(step, i, e)
        r[step+1] = r[step] + h * recovered_derivative(step, i)
        
    return s, e, i, r

s, e, i, r = seir_model()

def plot_me():
    s_plot = matplotlib.pyplot.plot(times, s, label = 'S')
    e_plot = matplotlib.pyplot.plot(times, e, label = 'E')
    i_plot = matplotlib.pyplot.plot(times, i, label = 'I')
    r_plot = matplotlib.pyplot.plot(times, r, label = 'R')
    matplotlib.pyplot.legend(('S', 'E', 'I', 'R'), loc = 'upper right')
    
    axes = matplotlib.pyplot.gca()
    axes.set_xlabel('Time in days')
    axes.set_ylabel('Number of persons')
    matplotlib.pyplot.xlim(xmin = 0.)
    matplotlib.pyplot.ylim(ymin = 0.)
    matplotlib.pyplot.show()

plot_me()


